{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ae4e29b4-454b-4f88-9b43-3409c080d2b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup OpenAI Agent\n",
    "import openai\n",
    "\n",
    "openai.api_key = \"sk-your-key\"\n",
    "from llama_index.agent import OpenAIAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c138547d-3e41-45f5-a9cc-afab22b56e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, List\n",
    "from llama_index.tools.tool_spec.base import BaseToolSpec\n",
    "import ast\n",
    "\n",
    "\n",
    "class PythonFileToolSpec(BaseToolSpec):\n",
    "    spec_functions = [\"function_definitions\", \"get_function\", \"get_functions\"]\n",
    "\n",
    "    def __init__(self, file_name: str) -> None:\n",
    "        f = open(file_name).read()\n",
    "        self.tree = ast.parse(f)\n",
    "\n",
    "    def function_definitions(self, external: Optional[bool] = True) -> str:\n",
    "        \"\"\"\n",
    "        Use this function to get the name and arguments of all function defintions in the python file\n",
    "        \"\"\"\n",
    "        functions = \"\"\n",
    "        for node in ast.walk(self.tree):\n",
    "            if isinstance(node, ast.FunctionDef):\n",
    "                if external and node.name.startswith(\"_\"):\n",
    "                    continue\n",
    "                functions += f\"\"\"\n",
    "name: {node.name}\n",
    "arguments: {ast.dump(node.args)}\n",
    "                    \"\"\"\n",
    "        return functions\n",
    "\n",
    "    def get_function(self, name: str) -> str:\n",
    "        for node in ast.walk(self.tree):\n",
    "            if isinstance(node, ast.FunctionDef):\n",
    "                if node.name == name:\n",
    "                    return f\"\"\"\n",
    "name: {node.name}\n",
    "arguments: {ast.dump(node.args)}\n",
    "docstring: {ast.get_docstring(node)}\n",
    "                        \"\"\"\n",
    "\n",
    "    def get_functions(self, names: List[str]) -> str:\n",
    "        functions = \"\"\n",
    "        for name in names:\n",
    "            functions += self.get_function(name) + \"\\n\"\n",
    "        return functions\n",
    "\n",
    "\n",
    "pyfile = PythonFileToolSpec(\"./numpy_linalg.py\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cd2a6e82-9c23-4ab1-89b6-1490cb1e1901",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the Agent with access to our tools\n",
    "agent = OpenAIAgent.from_tools(\n",
    "    pyfile.to_tool_list(),\n",
    "    system_prompt=\"\"\"\n",
    "    You are a specialized agent will help create descriptions of tools that are informative and instruct a user how to use the tools.\n",
    "    A tool is simply a function that will later be made callable by large language models like yourself.\n",
    "    Generally tool prompts will describe the function of the tool, and provide information on the arguments and return values.\n",
    "    \n",
    "    Here is an example of a function definition and a tool prompt:\n",
    "    def generate_image_variation(self, url: str, n: Optional[int] = 1, size: Optional[str] = '256x256') -> str:\n",
    "        ```Accepts the url of an image and uses OpenAIs api to generate a variation of the image.\n",
    "        This tool can take smaller images and create higher resolution variations, or vice versa.\n",
    "\n",
    "        When passing a url from \"generate_images\" ALWAYS pass the url exactly as it was returned from the function, including ALL query parameters\n",
    "        args:\n",
    "            url (str): The url of the image to create a variation of\n",
    "            n (Optional[int]): The number of images to generate. Defaults to 1.\n",
    "            size (Optional[str]): The size of the image(s) to generate. Defaults to 256x256. Other accepted values are 1024x1024 and 512x512\n",
    "        ```\n",
    "    \n",
    "    The first argument to all tools is self, as they are defined in a class\n",
    "\n",
    "    If you load a list of function names and arguments, you should help the user by creating tool prompts like the above.\n",
    "    \n",
    "    In general, the entire prompt should not be more than 5-10 lines, and should be short to save on tokens\n",
    "    \"\"\",\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ef7ff846-595d-4375-8a13-b4193171fec6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: get_functions with args: {\n",
      "  \"names\": [\"eig\", \"transpose\", \"solve\"]\n",
      "}\n",
      "Got output: \n",
      "name: eig\n",
      "arguments: arguments(posonlyargs=[], args=[arg(arg='a')], kwonlyargs=[], kw_defaults=[], defaults=[])\n",
      "docstring: Compute the eigenvalues and right eigenvectors of a square array.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "a : (..., M, M) array\n",
      "    Matrices for which the eigenvalues and right eigenvectors will\n",
      "    be computed\n",
      "\n",
      "Returns\n",
      "-------\n",
      "A namedtuple with the following attributes:\n",
      "\n",
      "eigenvalues : (..., M) array\n",
      "    The eigenvalues, each repeated according to its multiplicity.\n",
      "    The eigenvalues are not necessarily ordered. The resulting\n",
      "    array will be of complex type, unless the imaginary part is\n",
      "    zero in which case it will be cast to a real type. When `a`\n",
      "    is real the resulting eigenvalues will be real (0 imaginary\n",
      "    part) or occur in conjugate pairs\n",
      "\n",
      "eigenvectors : (..., M, M) array\n",
      "    The normalized (unit \"length\") eigenvectors, such that the\n",
      "    column ``eigenvectors[:,i]`` is the eigenvector corresponding to the\n",
      "    eigenvalue ``eigenvalues[i]``.\n",
      "\n",
      "Raises\n",
      "------\n",
      "LinAlgError\n",
      "    If the eigenvalue computation does not converge.\n",
      "\n",
      "See Also\n",
      "--------\n",
      "eigvals : eigenvalues of a non-symmetric array.\n",
      "eigh : eigenvalues and eigenvectors of a real symmetric or complex\n",
      "       Hermitian (conjugate symmetric) array.\n",
      "eigvalsh : eigenvalues of a real symmetric or complex Hermitian\n",
      "           (conjugate symmetric) array.\n",
      "scipy.linalg.eig : Similar function in SciPy that also solves the\n",
      "                   generalized eigenvalue problem.\n",
      "scipy.linalg.schur : Best choice for unitary and other non-Hermitian\n",
      "                     normal matrices.\n",
      "\n",
      "Notes\n",
      "-----\n",
      "\n",
      ".. versionadded:: 1.8.0\n",
      "\n",
      "Broadcasting rules apply, see the `numpy.linalg` documentation for\n",
      "details.\n",
      "\n",
      "This is implemented using the ``_geev`` LAPACK routines which compute\n",
      "the eigenvalues and eigenvectors of general square arrays.\n",
      "\n",
      "The number `w` is an eigenvalue of `a` if there exists a vector `v` such\n",
      "that ``a @ v = w * v``. Thus, the arrays `a`, `eigenvalues`, and\n",
      "`eigenvectors` satisfy the equations ``a @ eigenvectors[:,i] =\n",
      "eigenvalues[i] * eigenvalues[:,i]`` for :math:`i \\in \\{0,...,M-1\\}`.\n",
      "\n",
      "The array `eigenvectors` may not be of maximum rank, that is, some of the\n",
      "columns may be linearly dependent, although round-off error may obscure\n",
      "that fact. If the eigenvalues are all different, then theoretically the\n",
      "eigenvectors are linearly independent and `a` can be diagonalized by a\n",
      "similarity transformation using `eigenvectors`, i.e, ``inv(eigenvectors) @\n",
      "a @ eigenvectors`` is diagonal.\n",
      "\n",
      "For non-Hermitian normal matrices the SciPy function `scipy.linalg.schur`\n",
      "is preferred because the matrix `eigenvectors` is guaranteed to be\n",
      "unitary, which is not the case when using `eig`. The Schur factorization\n",
      "produces an upper triangular matrix rather than a diagonal matrix, but for\n",
      "normal matrices only the diagonal of the upper triangular matrix is\n",
      "needed, the rest is roundoff error.\n",
      "\n",
      "Finally, it is emphasized that `eigenvectors` consists of the *right* (as\n",
      "in right-hand side) eigenvectors of `a`. A vector `y` satisfying ``y.T @ a\n",
      "= z * y.T`` for some number `z` is called a *left* eigenvector of `a`,\n",
      "and, in general, the left and right eigenvectors of a matrix are not\n",
      "necessarily the (perhaps conjugate) transposes of each other.\n",
      "\n",
      "References\n",
      "----------\n",
      "G. Strang, *Linear Algebra and Its Applications*, 2nd Ed., Orlando, FL,\n",
      "Academic Press, Inc., 1980, Various pp.\n",
      "\n",
      "Examples\n",
      "--------\n",
      ">>> from numpy import linalg as LA\n",
      "\n",
      "(Almost) trivial example with real eigenvalues and eigenvectors.\n",
      "\n",
      ">>> eigenvalues, eigenvectors = LA.eig(np.diag((1, 2, 3)))\n",
      ">>> eigenvalues\n",
      "array([1., 2., 3.])\n",
      ">>> eigenvectors\n",
      "array([[1., 0., 0.],\n",
      "       [0., 1., 0.],\n",
      "       [0., 0., 1.]])\n",
      "\n",
      "Real matrix possessing complex eigenvalues and eigenvectors; note that the\n",
      "eigenvalues are complex conjugates of each other.\n",
      "\n",
      ">>> eigenvalues, eigenvectors = LA.eig(np.array([[1, -1], [1, 1]]))\n",
      ">>> eigenvalues\n",
      "array([1.+1.j, 1.-1.j])\n",
      ">>> eigenvectors\n",
      "array([[0.70710678+0.j        , 0.70710678-0.j        ],\n",
      "       [0.        -0.70710678j, 0.        +0.70710678j]])\n",
      "\n",
      "Complex-valued matrix with real eigenvalues (but complex-valued eigenvectors);\n",
      "note that ``a.conj().T == a``, i.e., `a` is Hermitian.\n",
      "\n",
      ">>> a = np.array([[1, 1j], [-1j, 1]])\n",
      ">>> eigenvalues, eigenvectors = LA.eig(a)\n",
      ">>> eigenvalues\n",
      "array([2.+0.j, 0.+0.j])\n",
      ">>> eigenvectors\n",
      "array([[ 0.        +0.70710678j,  0.70710678+0.j        ], # may vary\n",
      "       [ 0.70710678+0.j        , -0.        +0.70710678j]])\n",
      "\n",
      "Be careful about round-off error!\n",
      "\n",
      ">>> a = np.array([[1 + 1e-9, 0], [0, 1 - 1e-9]])\n",
      ">>> # Theor. eigenvalues are 1 +/- 1e-9\n",
      ">>> eigenvalues, eigenvectors = LA.eig(a)\n",
      ">>> eigenvalues\n",
      "array([1., 1.])\n",
      ">>> eigenvectors\n",
      "array([[1., 0.],\n",
      "       [0., 1.]])\n",
      "                        \n",
      "\n",
      "name: transpose\n",
      "arguments: arguments(posonlyargs=[], args=[arg(arg='a')], kwonlyargs=[], kw_defaults=[], defaults=[])\n",
      "docstring: Transpose each matrix in a stack of matrices.\n",
      "\n",
      "Unlike np.transpose, this only swaps the last two axes, rather than all of\n",
      "them\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "a : (...,M,N) array_like\n",
      "\n",
      "Returns\n",
      "-------\n",
      "aT : (...,N,M) ndarray\n",
      "                        \n",
      "\n",
      "name: solve\n",
      "arguments: arguments(posonlyargs=[], args=[arg(arg='a'), arg(arg='b')], kwonlyargs=[], kw_defaults=[], defaults=[])\n",
      "docstring: Solve a linear matrix equation, or system of linear scalar equations.\n",
      "\n",
      "Computes the \"exact\" solution, `x`, of the well-determined, i.e., full\n",
      "rank, linear matrix equation `ax = b`.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "a : (..., M, M) array_like\n",
      "    Coefficient matrix.\n",
      "b : {(..., M,), (..., M, K)}, array_like\n",
      "    Ordinate or \"dependent variable\" values.\n",
      "\n",
      "Returns\n",
      "-------\n",
      "x : {(..., M,), (..., M, K)} ndarray\n",
      "    Solution to the system a x = b.  Returned shape is identical to `b`.\n",
      "\n",
      "Raises\n",
      "------\n",
      "LinAlgError\n",
      "    If `a` is singular or not square.\n",
      "\n",
      "See Also\n",
      "--------\n",
      "scipy.linalg.solve : Similar function in SciPy.\n",
      "\n",
      "Notes\n",
      "-----\n",
      "\n",
      ".. versionadded:: 1.8.0\n",
      "\n",
      "Broadcasting rules apply, see the `numpy.linalg` documentation for\n",
      "details.\n",
      "\n",
      "The solutions are computed using LAPACK routine ``_gesv``.\n",
      "\n",
      "`a` must be square and of full-rank, i.e., all rows (or, equivalently,\n",
      "columns) must be linearly independent; if either is not true, use\n",
      "`lstsq` for the least-squares best \"solution\" of the\n",
      "system/equation.\n",
      "\n",
      "References\n",
      "----------\n",
      ".. [1] G. Strang, *Linear Algebra and Its Applications*, 2nd Ed., Orlando,\n",
      "       FL, Academic Press, Inc., 1980, pg. 22.\n",
      "\n",
      "Examples\n",
      "--------\n",
      "Solve the system of equations ``x0 + 2 * x1 = 1`` and ``3 * x0 + 5 * x1 = 2``:\n",
      "\n",
      ">>> a = np.array([[1, 2], [3, 5]])\n",
      ">>> b = np.array([1, 2])\n",
      ">>> x = np.linalg.solve(a, b)\n",
      ">>> x\n",
      "array([-1.,  1.])\n",
      "\n",
      "Check that the solution is correct:\n",
      "\n",
      ">>> np.allclose(np.dot(a, x), b)\n",
      "True\n",
      "                        \n",
      "\n",
      "========================\n",
      "Here is a function definition using only built-in Python types for the `eig` function:\n",
      "\n",
      "```python\n",
      "def compute_eigenvalues_and_eigenvectors(a: List[List[float]]) -> Tuple[List[complex], List[List[complex]]]:\n",
      "    \"\"\"\n",
      "    Compute the eigenvalues and right eigenvectors of a square array.\n",
      "\n",
      "    Parameters:\n",
      "        a (List[List[float]]): Matrices for which the eigenvalues and right eigenvectors will be computed\n",
      "\n",
      "    Returns:\n",
      "        Tuple[List[complex], List[List[complex]]]: A tuple containing the eigenvalues and eigenvectors.\n",
      "            The eigenvalues are a list of complex numbers, each repeated according to its multiplicity.\n",
      "            The eigenvectors are a list of lists, where each inner list represents a normalized eigenvector.\n",
      "            The column eigenvectors[i] is the eigenvector corresponding to the eigenvalue eigenvalues[i].\n",
      "    \"\"\"\n",
      "    pass\n",
      "```\n",
      "\n",
      "And here is a function definition using only built-in Python types for the `transpose` function:\n",
      "\n",
      "```python\n",
      "def transpose_matrix(a: List[List[float]]) -> List[List[float]]:\n",
      "    \"\"\"\n",
      "    Transpose each matrix in a stack of matrices.\n",
      "\n",
      "    Parameters:\n",
      "        a (List[List[float]]): The matrix to transpose\n",
      "\n",
      "    Returns:\n",
      "        List[List[float]]: The transposed matrix\n",
      "    \"\"\"\n",
      "    pass\n",
      "```\n",
      "\n",
      "And finally, here is a function definition using only built-in Python types for the `solve` function:\n",
      "\n",
      "```python\n",
      "def solve_linear_equation(a: List[List[float]], b: Union[List[float], List[List[float]]]) -> Union[List[float], List[List[float]]]:\n",
      "    \"\"\"\n",
      "    Solve a linear matrix equation, or system of linear scalar equations.\n",
      "\n",
      "    Parameters:\n",
      "        a (List[List[float]]): Coefficient matrix.\n",
      "        b (Union[List[float], List[List[float]]]): Ordinate or \"dependent variable\" values.\n",
      "\n",
      "    Returns:\n",
      "        Union[List[float], List[List[float]]]: Solution to the system a x = b.\n",
      "            The returned shape is identical to b.\n",
      "    \"\"\"\n",
      "    pass\n",
      "```\n",
      "\n",
      "These function definitions use the `List` and `Tuple` types from the built-in `typing` module to represent lists and tuples respectively. The `float` type is used to represent floating-point numbers. The `Union` type is used to indicate that the function can return either a single value or a list of values.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    agent.chat(\n",
    "        \"\"\"\n",
    "Load the eig, transpose and solve functions from the python file,\n",
    "and then write a function defintion using ONLY! builtin python types (List, float, Tuple)\n",
    "with a short 5-10 line doc string tool prompts for the functions that only has a small description and arguments\n",
    "\"\"\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4f82c1a0-6c1e-4efb-99f0-1e850837b6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Numpy tool spec.\"\"\"\n",
    "\n",
    "from llama_index.tools.tool_spec.base import BaseToolSpec\n",
    "from typing import Optional, List, Tuple, Union\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class NumpyToolSpec(BaseToolSpec):\n",
    "    \"\"\"Numpy Tool Spec\"\"\"\n",
    "\n",
    "    spec_functions = [\n",
    "        \"compute_eigenvalues_and_eigenvectors\",\n",
    "        \"transpose_matrix\",\n",
    "        \"solve_linear_equation\",\n",
    "    ]\n",
    "\n",
    "    def compute_eigenvalues_and_eigenvectors(\n",
    "        self, a: List[List[float]]\n",
    "    ) -> Tuple[List[complex], List[List[complex]]]:\n",
    "        \"\"\"\n",
    "        Compute the eigenvalues and right eigenvectors of a square array.\n",
    "\n",
    "        Parameters:\n",
    "            a (List[List[float]]): Matrices for which the eigenvalues and right eigenvectors will be computed\n",
    "\n",
    "        Returns:\n",
    "            Tuple[List[complex], List[List[complex]]]: A tuple containing the eigenvalues and eigenvectors.\n",
    "                The eigenvalues are a list of complex numbers, each repeated according to its multiplicity.\n",
    "                The eigenvectors are a list of lists, where each inner list represents a normalized eigenvector.\n",
    "                The column eigenvectors[i] is the eigenvector corresponding to the eigenvalue eigenvalues[i].\n",
    "        \"\"\"\n",
    "        return np.linalg.eig(a)\n",
    "\n",
    "    def transpose_matrix(self, a: List[List[float]]) -> List[List[float]]:\n",
    "        \"\"\"\n",
    "        Transpose each matrix in a stack of matrices.\n",
    "\n",
    "        Parameters:\n",
    "            a (List[List[float]]): The matrix to transpose\n",
    "\n",
    "        Returns:\n",
    "            List[List[float]]: The transposed matrix\n",
    "        \"\"\"\n",
    "        return np.transpose(a)\n",
    "\n",
    "    def solve_linear_equation(\n",
    "        self, a: List[List[float]], b: Union[List[float], List[List[float]]]\n",
    "    ) -> Union[List[float], List[List[float]]]:\n",
    "        \"\"\"\n",
    "        Solve a linear matrix equation, or system of linear scalar equations.\n",
    "\n",
    "        Parameters:\n",
    "            a (List[List[float]]): Coefficient matrix.\n",
    "            b (Union[List[float], List[List[float]]]): Ordinate or \"dependent variable\" values.\n",
    "\n",
    "        Returns:\n",
    "            Union[List[float], List[List[float]]]: Solution to the system a x = b.\n",
    "                The returned shape is identical to b.\n",
    "        \"\"\"\n",
    "        return np.linalg.solve(a, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1bd6aeed-9c92-417b-a0f7-942fb2bcf79d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import OpenAIAgent\n",
    "\n",
    "agent = OpenAIAgent.from_tools(NumpyToolSpec().to_tool_list(), verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cca60dad-bfd9-4562-9bf5-01bb871ffdfa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: solve_linear_equation with args: {\n",
      "  \"a\": [[1, 2], [3, 5]],\n",
      "  \"b\": [1, 2]\n",
      "}\n",
      "Got output: [-1.  1.]\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: transpose_matrix with args: {\n",
      "  \"a\": [[1, 2], [3, 5]]\n",
      "}\n",
      "Got output: [[1 3]\n",
      " [2 5]]\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: compute_eigenvalues_and_eigenvectors with args: {\n",
      "  \"a\": [[1, 3], [2, 5]]\n",
      "}\n",
      "Got output: (array([-0.16227766,  6.16227766]), array([[-0.93246475, -0.50245469],\n",
      "       [ 0.36126098, -0.86460354]]))\n",
      "========================\n",
      "The solution to the system of equations is x0 = -1 and x1 = 1.\n",
      "\n",
      "The transpose of the coefficient matrix is [[1, 3], [2, 5]].\n",
      "\n",
      "The eigenvalues of the coefficient matrix are -0.16227766 and 6.16227766. The corresponding eigenvectors are [-0.93246475, -0.50245469] and [0.36126098, -0.86460354].\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    agent.chat(\n",
    "        \"\"\"\n",
    "Using the tools provided, solve the system of equations ``x0 + 2 * x1 = 1`` and ``3 * x0 + 5 * x1 = 2``, then transpose the coefficent matrix and compute the eigenvalues\n",
    "\"\"\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37cf4296-f558-4127-bace-9bb9374192aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
